/*=============================================================
* 	Copyright (C) 2024 USTC All rights reserved.
* 	
* 	filename : core.cpp
* 	Author: Wei Jinsong
* 	Mail: vitus@mainl.ustc.edu.cn
* 	Date 2024-05-13
* 	Abstract 
* 		
*
=============================================================*/

#include <random>
#include <iostream>
#include "core.h"
#include "utils.h"

using namespace sc_core;
using namespace std;

BaseCore::BaseCore(sc_module_name name, uint32_t id, sc_clock& clock, string file)
            // , uint32_t recv_pipeline, uint32_t send_pipeline)
    : sc_core::sc_module(name)
    , mID(id)
    , spikeFifo(20480)
    , recvEvent((string(name)+"_recvEvent").c_str())
    // , recvPipeline(recv_pipeline), sendPipeline(send_pipeline)
    , recvSpikePeq((string(name) + "_recvSpikePeq").c_str())
    , sendSpikePeq((string(name) + "_sendSpikePeq").c_str())
{
    clk(clock);
    SC_TRACEX(clock);
    inSpikeSig.valid = 0;
    outSpikeSig.valid = 0;
    sc_trace_spike("inSpike", inSpikeSig);
    sc_trace_spike("outSpike", outSpikeSig);

    clkPeriod = clock.period();
    type = "null";

    SC_THREAD(exec);
    sensitive << clk.posedge_event();
    dont_initialize();

    SC_THREAD(fs_exec);
    sensitive << clk.posedge_event();
    dont_initialize();
    
    SC_THREAD(recv_spike);
    sensitive << recvSpikePeq.get_event();
    dont_initialize();

    SC_THREAD(generate_spike);
    sensitive << sendSpikePeq.get_event();
    dont_initialize();
}

BaseCore::~BaseCore() 
{
}

void
BaseCore::sc_trace_spike(string name, SpikePayload& spike)
{
    string spath = this->name();
    spath = spath+"."+name;
    SC_TRACE(spike.srcid   , spath+".srcid"   );
    SC_TRACE(spike.srcneuid, spath+".srcneuid");
    SC_TRACE(spike.tgtid   , spath+".tgtid"   );
    SC_TRACE(spike.tgtneuid, spath+".tgtneuid");
    SC_TRACE(spike.step    , spath+".step"    );
    SC_TRACE(spike.channel , spath+".channel" );
    SC_TRACE(spike.height  , spath+".height"  );
    SC_TRACE(spike.width   , spath+".width"   );
    SC_TRACE(spike.phase   , spath+".phase"   );
    SC_TRACE(spike.switch_phase, spath+".switch_phase");
    SC_TRACE(spike.valid   , spath+".valid"   );
}

void 
BaseCore::set_featureMap(string file)
{
    ifstream j_if(file);
    if (!j_if.is_open()) {
        return ;
    }
    json j;
    j_if >> j;

    featureType = j["type"];
    featureTH   = j["th"];
    featureMaps.insert(featureMaps.begin(), j["features"].begin(), j["features"].end());
    featureToken= 0;
}

void
BaseCore::exec_input()
{
    random_device e;
    uniform_real_distribution<float> u(0,1);
    // for (;i<featureMaps.size(); i++) {
    while (true) {
        int i=featureToken;
        featureToken = featureToken + 1;
        if (featureToken>featureMaps.size()) {
            featureToken = 0;
            break;
        }
        bool spike = false;
        switch (featureType) {
            case 0 : 
                spike = featureMaps[i]>=featureTH;
                featureMaps[i] = (spike) ? featureMaps[i]-featureTH : featureTH;
                break;
            case 1 :
                spike = featureMaps[i]>=u(e);
                break;
            default :
                break;
        }

        if (spike) {
            SpikePayload spike;
            spike.srcid = mID;
            spike.srcneuid = i;
            vector<SpikePayload*>new_spikes = synapse->get_axon(&spike);

            for (int s=0; s<new_spikes.size(); s++) {
                sendSpikePeq.notify(*new_spikes[s], sendPipeline*clkPeriod);
            }
            break;
        }
    }
}

bool
BaseCore::spike_transport_fw(SpikePayload &trans, sc_core::sc_time &delay)
{
    recvSpikePeq.notify(trans, recvPipeline*clkPeriod);
    return true;
}

void
BaseCore::recv_spike()
{
    SpikePayload* spike = NULL;
    while (true) {
        while (spikeFifo.num_free()>0 
                && (spike = recvSpikePeq.get_next_transaction())!=NULL) {
            inSpikeSig = *spike;
            spikeFifo.nb_write(spike);

            if (spike->switch_phase==false) {
                _inSpikeCnt[spike->tgtneuid]+=1;
            }
            wait(clk.posedge_event());
            inSpikeSig.valid = false;
        }
        if (spikeFifo.num_free()<=0){
            cout << sc_time_stamp() << "core" << mID << " recv spike fifo full!" << endl;
        }
        wait();
    }
}

void 
BaseCore::generate_spike()
{
    SpikePayload* spike = NULL;
    while (true) {
        while ( (spike = sendSpikePeq.get_next_transaction())!=NULL ) {
            sc_time delay = clkPeriod;
            outSpikeSig = *spike;

            if (spike->switch_phase==false) {
                _outSpikeCnt[spike->tgtneuid]+=1;
            }
            // if (_outSpikeCnt.find(spike->tgtneuid) == _outSpikeCnt.end()) {
            //     _outSpikeCnt[spike->tgtneuid]=1;
            // }else {
            //     _outSpikeCnt[spike->tgtneuid]+=1;
            // }
            txSpikePort->spike_transport_fw(*spike, delay);

            wait(clk.posedge_event());
            outSpikeSig.valid = false;
        }
        wait();
    }
}

int64_t 
BaseCore::get_recv_spike() 
{
    uint64_t num = 0;
    for (auto cnt : _inSpikeCnt) {
        num = num+cnt.second;
    }
    return num;
}

int64_t 
BaseCore::get_send_spike() 
{
    uint64_t num = 0;
    for (auto cnt : _outSpikeCnt) {
        num = num+cnt.second;
    }
    return num;
}

void 
BaseCore::print_spike_sta(ostream &os)
{
    if (this->type=="output") 
        if (_inSpikeCnt.size()!=0 || _outSpikeCnt.size()!=0) {
            os << "Core " << mID << " " << type << endl;
            os << "\t" << "Receive Spike :" << endl
               << "\t\t" ;
            for (auto cnt : _inSpikeCnt) {
                os << "(" <<  cnt.first << "," << cnt.second << ")" ;
            }
            os << endl;
            os << "\t" << "Send Spike :" << endl
               << "\t\t" ;
            for (auto cnt : _outSpikeCnt) {
                os << "(" <<  cnt.first << "," << cnt.second << ")" ;
            }
            os << endl;
        } 
}

/*
{
"neuron_num": 10,
"recv_pipeline":1,
"send_pipeline":1,
"threshold" : 10.0,
"tua" : 2.0
}
*/
void
SpikeNeuronCore::set_params(string file)
{
    ifstream j_if(file);
    if (!j_if.is_open()) {
        neuNum = 0;
        return ;
    }
    json j;
    j_if >> j;

    neuNum = j["neuron_num"];
    recvPipeline = j["recv_pipeline"];
    sendPipeline = j["send_pipeline"];
    _th = j["threshold"];
    _t  = j["tua"];

    // _Vmem.resize(neuNum);
    _Spikes = new sc_signal<bool>[neuNum];
    _Vmem = new sc_signal<double>[neuNum];

    for (int i=0; i<neuNum; i++) {
        SC_TRACE(_Vmem[i],string(this->name())+"._Vmem"+to_string(i));
        SC_TRACE(_Spikes[i],string(this->name())+"._Spikes"+to_string(i));
    }

    j_if.close();
}

SpikeNeuronCore::SpikeNeuronCore(string name, uint32_t id, sc_clock& clock
        , string file, string infile)
    : BaseCore(name.c_str(), id, clock, file)
{
    set_params(file);
    set_featureMap(infile);
}

SpikeNeuronCore::~SpikeNeuronCore()
{
}

void 
SpikeNeuronCore::exec()
{
    SpikePayload* spike = NULL;
    while (true) {
        exec_input();
        for (int i=0; i<neuNum; i++) {
            if (_Vmem[i]>_th) {
                SpikePayload spike;
                spike.srcid = mID;
                spike.srcneuid = i;
                vector<SpikePayload*>new_spikes = synapse->get_axon(&spike);

                for (int s=0; s<new_spikes.size(); s++) {
                    sendSpikePeq.notify(*new_spikes[s], sendPipeline*clkPeriod);
                }
                _Vmem[i] = 0;
            }
        }
        
        if (spikeFifo.nb_read(spike)) {
            vector<SpikePayload*> weights = synapse->get_dendrite(spike);

            for (int i=0; i<weights.size(); i++) {
                uint32_t tgtneuid = weights[i]->tgtneuid;
                double* weight = (double*)weights[i]->weights;
                _Vmem[tgtneuid] = _Vmem[tgtneuid] + *weight;
            }
            delete spike;
            spike = NULL;
        }
        wait();
    }
}

FSNeuronCore::FSNeuronCore(string name, uint32_t id, sc_clock& clock
        , string file, string infile)
    : BaseCore(name.c_str(), id, clock, file)
{
    _currStep = 0;
    _currPhase = 0;
    _currVmem = 0;
    fs_phase = 0;
    _fs_exec_state = false;
    _fs_recv_state = true;
    _switchReady = false;

    set_params(file);
    set_featureMap(infile);

    SC_TRACEX(_currPhase);
    SC_TRACEX(_currStep);
    SC_TRACEX(_currVmem);
    SC_TRACEX(_switchReady);
    SC_TRACEX(_fs_exec_state);
    SC_TRACEX(_fs_recv_state);
}

FSNeuronCore::~FSNeuronCore()
{
}


void 
FSNeuronCore::set_featureMap(string file)
{
    ifstream j_if(file);
    if (!j_if.is_open()) {
        return ;
    }
    LOG(file);
    json j;
    j_if >> j;
    type = "input";
    _featureParam.channel = j["channel"];
    _featureParam.height = j["height"];
    _featureParam.width = j["width"];

    neuNum = _featureParam.channel*_featureParam.height*_featureParam.width;
    featureMaps.insert(featureMaps.begin(), j["features"].begin(), j["features"].end());
    featureToken= 0;

    recvPipeline = j["recv_pipeline"];
    sendPipeline = j["send_pipeline"];
    _Vmem[0].insert(_Vmem[0].begin(), j["features"].begin(), j["features"].end());
    _Vmem[1].resize(neuNum, 0);

    fs_phase = j["fs_phase"];
    fs_k = j["fs_k"];
    fs_h.insert(fs_h.begin(),j["fs_h"].begin(), j["fs_h"].end());
    fs_d.insert(fs_d.begin(),j["fs_d"].begin(), j["fs_d"].end());
    fs_t.insert(fs_t.begin(),j["fs_t"].begin(), j["fs_t"].end());

    LOG(neuNum);
    LOG(_Vmem[0].size());
    LOG(_Vmem[1].size());
    for (int i=0; i<neuNum; i++) {
        SC_TRACE(_Vmem[0][i],string(this->name())+"._Vmem0_"+to_string(i));
        SC_TRACE(_Vmem[1][i],string(this->name())+"._Vmem1_"+to_string(i));
    }
    _fs_exec_state = true;
    _fs_recv_state = false;
}

void
FSNeuronCore::set_params(string file)
{
    ifstream j_if(file);
    if (!j_if.is_open()) {
        neuNum = 0;
        return;
    }
    LOG(file);
    json j;
    j_if >> j;

    
    type = j["type"];
    if (j.find("label")!=j.end()) {
        label = j["label"];
    }
    if (type=="conv") {
        _poolParam.kh = j["pool_h"];
        _poolParam.kw = j["pool_w"];
        _poolParam.sh = j["pool_s"];
        _poolParam.sw = j["pool_s"];

        _convParam.ic = j["inChannel"];
        _convParam.oc = j["conv_c"];
        _convParam.kh = j["conv_h"];
        _convParam.kw = j["conv_w"];
        _convParam.sh = j["conv_s"];
        _convParam.sw = j["conv_s"];

        _featureParam.channel = j["channel"];
        _featureParam.height = j["height"];
        _featureParam.width = j["width"];

        neuNum = _featureParam.channel*_featureParam.height*_featureParam.width;
    }else if (type=="linear") {
        _convParam.ic = j["inChannel"];
        _convParam.oc = j["channel"];
        _poolParam.kh = j["pool_h"];
        _poolParam.kw = j["pool_w"];
        _poolParam.sh = j["pool_s"];
        _poolParam.sw = j["pool_s"];
        _featureParam.channel = j["channel"];
        neuNum = _featureParam.channel;
    }else {
        neuNum = j["channel"];
    }
    if (j.find("bias")!=j.end()) {
        biasParam.insert(biasParam.begin(), j["bias"].begin(), j["bias"].end());
    }else{
    }

    fs_phase = j["fs_phase"];
    fs_k = j["fs_k"];
    fs_h.insert(fs_h.begin(),j["fs_h"].begin(), j["fs_h"].end());
    fs_d.insert(fs_d.begin(),j["fs_d"].begin(), j["fs_d"].end());
    fs_t.insert(fs_t.begin(),j["fs_t"].begin(), j["fs_t"].end());

    if (j.find("fs_tau")!=j.end()) {
        fs_tau.insert(fs_tau.begin(),j["fs_tau"].begin(), j["fs_tau"].end());
    }

    recvPipeline = j["recv_pipeline"];
    sendPipeline = j["send_pipeline"];
    _Vmem[0].resize(neuNum,0);
    _Vmem[1].resize(neuNum,0);

    for (int i=0; i<neuNum; i++) {
        // SC_TRACE(_Vmem[0][i],string(this->name())+"._Vmem0_"+to_string(i));
        // SC_TRACE(_Vmem[1][i],string(this->name())+"._Vmem1_"+to_string(i));
    }

    j_if.close();
}

void 
FSNeuronCore::send_spike(SpikePayload &spike)
{
    vector<SpikePayload*> new_spikes = synapse->get_axon(&spike);
    for (int s=0; s<new_spikes.size(); s++) {
        sendSpikePeq.notify(*new_spikes[s], sendPipeline*clkPeriod);
        wait(clk.posedge_event());
    }
}

void 
FSNeuronCore::send_spike(int32_t srcneuid, int32_t phase, int32_t step, bool switch_phase, double value)
{
    SpikePayload spike;
    spike.srcid = mID;
    spike.srcneuid = srcneuid;
    spike.step = step;
    spike.phase = phase;
    spike.switch_phase = switch_phase;
    spike.value = value;
    send_spike(spike);
}

bool 
FSNeuronCore::update_neuron(int32_t currVmem, int32_t currStep, int32_t neuid)
{
    if (_Vmem[currVmem][neuid] >= fs_t[currStep]) {
        _Vmem[currVmem][neuid] -= fs_h[_currStep];
        return true;
    }
    return false;
}

void 
FSNeuronCore::exec_input()
{
    if (_currPhase>=fs_phase) {
        return;
    }
    LOG(sc_time_stamp());
    LOG(type);
    LOG(_currStep);
    LOG(_currPhase);
    LOG(_currVmem);
    for (int i=0; i<neuNum; i++) {
        if (update_neuron(_currVmem, _currStep, i)) {
            send_spike(i, _currPhase, _currStep, false);
        }
    }
    _currStep=_currStep+1;
    if (_currStep>=fs_k) {
        _currStep = 0;
        _currPhase=_currPhase + 1;
        _fs_exec_state = false;
        if (this->type!="output")
            send_spike(0, _currPhase, 0, true);
    }
}

void 
FSNeuronCore::exec_vmem() 
{
    if (_currPhase>=fs_phase) {
        return;
    }
    for (int i=0; i<neuNum; i++) {
        if (_fs_exec_state && update_neuron(_currVmem, _currStep, i)) {
            send_spike(i, _currPhase, _currStep, false);
        }
    }
    _currStep=_currStep+1;
    if (_currStep>=fs_k) {
        _currStep = 0;
        _currPhase=_currPhase + 1;
        _fs_exec_state = false;
        if (this->type!="output")
            send_spike(0, _currPhase, 0, true);
    }
}

void 
FSNeuronCore::fs_exec() 
{
    while (true) {
        if (_fs_exec_state){
            if (type=="input") {
                exec_input();
            } else if (type!="output" && type!="null") {
                exec_vmem();
            }
        }else if (_switchReady) {
            if (_currPhase<fs_phase) {
                _fs_exec_state = true;
            }else {
                _fs_exec_state = false;
            }
            wait();
        }
        wait();
    }
}

void 
FSNeuronCore::output_sta(SpikePayload *spike)
{
    _Vmem[1-_currVmem][spike->tgtneuid] += fs_d[spike->step];
}

void 
FSNeuronCore::exec()
{
    SpikePayload* spike = NULL;
    while (true) {
        if (_fs_recv_state) {
            if (spikeFifo.nb_read(spike)) {
                if (this->type=="output") {
                    if (spike->switch_phase && spike->phase==fs_phase-1) {
                        cout << "Runing Time : " << sc_time_stamp() << endl;
                        sc_stop();
                    }else if (spike->switch_phase==false) {
                        output_sta(spike);
                    }
                }else if (spike->switch_phase) {
                    _fs_recv_state = false;
                } else {
                    vector<SpikePayload*> weights = synapse->get_dendrite(spike);
                    int step = spike->step;
                    for (int i=0; i<weights.size(); i++) {

                        double* weight = (double*)weights[i]->weights;
                        uint32_t tgtneuid = 0;

                        if (this->type=="conv") {
                            if (weights[i]->height>=_featureParam.height || weights[i]->width>=_featureParam.width) {
                                continue;
                            }
                            tgtneuid = weights[i]->channel*(_featureParam.height*_featureParam.width)
                                     + weights[i]->height*_featureParam.width
                                     + weights[i]->width;
                        }else if (this->type=="linear") {
                            tgtneuid = weights[i]->channel;
                        }
                        _Vmem[1-_currVmem][tgtneuid] +=  spike->value*( *weight ) * fs_d[step] / (_poolParam.kh*_poolParam.kw);
                    }
                }
                delete spike;
                spike = NULL;
            }
        }else {
            if (_switchReady) {
                _fs_recv_state = this->type!="input";
                _currVmem = 1-_currVmem;
            }
            wait();
        }
        _switchReady = (_fs_exec_state==false) && (_fs_recv_state==false);
        wait();
    }
}

void 
FSNeuronCore::print_spike_sta(ostream &os)
{
    BaseCore::print_spike_sta(os);
    if (_inSpikeCnt.size()!=0 || _outSpikeCnt.size()!=0) {
        if (this->type=="output") {
            os << "\t Vmem[0]" << endl;
            os << "\t\t";
            for (int i=0; i<_Vmem[0].size(); i++) {
                os << _Vmem[0][i] << ", ";
            }
            os << endl;
            os << "\t Vmem[1]" <<endl;
            os << "\t\t";
            for (int i=0; i<_Vmem[1].size(); i++) {
                os << _Vmem[1][i] << ", ";
            }
            os << endl;
        }
    } 
}

AFSNeuronCore::AFSNeuronCore(string name, uint32_t mId, sc_clock& clock
            , string file, string infile)
    :FSNeuronCore(name, mId, clock, file, infile)
{
    if (neuNum!=0) {
        _ExecVmem.resize(neuNum, 0);
    }
}
AFSNeuronCore::~AFSNeuronCore()
{
}

void 
AFSNeuronCore::exec_input() 
{
    if (_currPhase>=fs_phase) {
        return;
    }
    LOG(sc_time_stamp());
    LOG(type);
    LOG(_currStep);
    LOG(_currPhase);
    LOG(_currVmem);
    for (int i=0; i<neuNum; i++) {
        int featureid = _currPhase*neuNum + i;
        if (featureMaps[featureid]>=1.0) {
            send_spike(i, _currPhase, _currStep, false, featureMaps[featureid]);
        }
    }

    _currStep=_currStep+1;
    if (_currStep>=fs_k) {
        _currStep = 0;
        _currPhase= _currPhase+1;
        _fs_exec_state = false;
        if (this->type!="output")
            send_spike(0, _currPhase, 0, true);
    }
    return ;
}


bool 
AFSNeuronCore::update_neuron(int32_t currVmem, int32_t currStep, int32_t neuid)
{
    int channel;
    double bias = 0;
    if (this->type=="conv") {
        channel = neuid / (_featureParam.height*_featureParam.width);
    }else {
        channel = neuid;
    }
    if (biasParam.size()>channel) {
        bias = biasParam[channel];
    }
    _ExecVmem[neuid] += (_Vmem[currVmem][neuid] + bias - (_ExecVmem[neuid]-fs_h[0]))/fs_tau[0];
    _Vmem[currVmem][neuid] = 0.0;

    if (_ExecVmem[neuid]>=fs_t[0]) {
        _ExecVmem[neuid] = fs_h[0];
        return true;
    }
    return false;
}

void 
AFSNeuronCore::output_sta(SpikePayload *spike)
{
    _dvs_output[spike->tgtneuid/10] += 1;
    _dvs_output2[spike->tgtneuid%11] += 1;
}

void 
AFSNeuronCore::print_spike_sta(ostream& os)
{
    FSNeuronCore::print_spike_sta(os);
    if (_dvs_output.size()!=0) {
        os << "\tlabel " << label << endl;
        if (this->type=="output") {
            os << "\t dvs_output" <<endl;
            os << "\t\t";
            for (auto cnt : _dvs_output) {
                os << "(" <<  cnt.first << "," << cnt.second << ")" ;
            }
            os << endl;
        }
        if (this->type=="output") {
            os << "\t dvs_output2" <<endl;
            os << "\t\t";
            for (auto cnt : _dvs_output2) {
                os << "(" <<  cnt.first << "," << cnt.second << ")" ;
            }
            os << endl;
        }
    }
}

